package tools;
import tools.LRParser;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

public class IRGenerator {
    private int tempCount = 0;
    private int labelCount = 0;

    public List<String> generate(LRParser.Node node) {
        List<String> code = new ArrayList<>();
        generateFromNode(node, code);
        return code;
    }

    private void generateFromNode(LRParser.Node node, List<String> code) {
        if (node == null) return;
        if (node instanceof LRParser.NonTerminalNode nt) {
            if (nt.getChildren().isEmpty()) return;

            switch (nt.getSymbol()) {
                case "statement_list" -> nt.getChildren().forEach(child -> generateFromNode(child, code));
                case "statement" -> nt.getChildren().forEach(child -> generateFromNode(child, code));
                case "declaration" -> handleDeclaration(nt, code);
                case "assignment" -> handleAssignment(nt, code);
                case "output" -> handleOutput(nt, code);
                case "selection" -> handleIfElse(nt, code);
                case "loop" -> handleWhile(nt, code);
                default -> nt.getChildren().forEach(child -> generateFromNode(child, code));
            }
        }
    }

    private void handleDeclaration(LRParser.NonTerminalNode decl, List<String> code) {
        String id = getTokenValue(decl.getChildren().get(1));
        if (decl.getChildren().size() == 3) {
            code.add("int " + id);
        } else {
            LRParser.NonTerminalNode expr = (LRParser.NonTerminalNode) decl.getChildren().get(3);
            String val = generateExpression(expr, code);
            code.add("int "+id + " = " + val);
        }
    }


    //a = b + 3 * (2 - 1);
    private void handleAssignment(LRParser.NonTerminalNode node, List<String> code) {
        String id = getTokenValue(node.getChildren().get(0));
        //等于三代表只是简单的赋值语句
        if (node.getChildren().size() == 3) {
            String op = getTokenValue(node.getChildren().get(1));
            String temp = newTemp();
            if ("++".equals(op)) {
                code.add(temp + " = " + id + " + 1");
                code.add(id + " = " + temp);
            } else if ("--".equals(op)) {
                code.add(temp + " = " + id + " - 1");
                code.add(id + " = " + temp);
            }
        } else {

            //先计算表达式右边
            LRParser.NonTerminalNode expr = (LRParser.NonTerminalNode) node.getChildren().get(2);
            //开始计算表达式右边
            //System.out.println("开始计算表达式右边");
            String val = generateExpression(expr, code);
            code.add(id + " = " + val);
        }
    }

    private void handleOutput(LRParser.NonTerminalNode node, List<String> code) {
        LRParser.NonTerminalNode expr = (LRParser.NonTerminalNode) node.getChildren().get(2);
        String val = generateExpression(expr, code);
        code.add("printf " + val);
    }

    private void handleIfElse(LRParser.NonTerminalNode node, List<String> code) {
        String trueLabel = newLabel();
        String falseLabel = newLabel();
        String endLabel = newLabel();

        LRParser.NonTerminalNode cond = (LRParser.NonTerminalNode) node.getChildren().get(2);
        generateCondition(cond, code, trueLabel, falseLabel);

        code.add(trueLabel + ":");
        generateFromNode(node.getChildren().get(5), code);
        code.add("goto " + endLabel);

        code.add(falseLabel + ":");
        generateFromNode(node.getChildren().get(9), code);

        code.add(endLabel + ":");
    }

    private void handleWhile(LRParser.NonTerminalNode node, List<String> code) {
        String startLabel = newLabel();
        String bodyLabel = newLabel();
        String endLabel = newLabel();

        code.add(startLabel + ":");
        LRParser.NonTerminalNode cond = (LRParser.NonTerminalNode) node.getChildren().get(2);
        generateCondition(cond, code, bodyLabel, endLabel);

        code.add(bodyLabel + ":");
        generateFromNode(node.getChildren().get(5), code);
        code.add("goto " + startLabel);
        code.add(endLabel + ":");
    }

    private void generateCondition(LRParser.NonTerminalNode cond, List<String> code, String trueLabel, String falseLabel) {
        String left = generateExpression((LRParser.NonTerminalNode) cond.getChildren().get(0), code);
        String op = getTokenValue(((LRParser.NonTerminalNode) cond.getChildren().get(1)).getChildren().get(0));
        String right = generateExpression((LRParser.NonTerminalNode) cond.getChildren().get(2), code);

        code.add("if " + left + " " + op + " " + right + " goto " + trueLabel);
        code.add("goto " + falseLabel);
    }

    private String generateExpression(LRParser.NonTerminalNode node, List<String> code) {
        //进来的默认是一个term＋expression_tail的结构
        String left = generateTerm((LRParser.NonTerminalNode) node.getChildren().get(0), code);
        //后面还有，就调用处理expression_tail
       // System.out.println("处理最左边term："+left);
        if (node.getChildren().size() > 1) {
            LRParser.NonTerminalNode exprTail = (LRParser.NonTerminalNode) node.getChildren().get(1);
            return handleAdditiveExpression(left, exprTail, code);
        }
        return left;
    }

    private String handleAdditiveExpression(String left, LRParser.NonTerminalNode tail, List<String> code) {
        if (tail.getChildren().isEmpty()) return left;

        String op = getTokenValue(tail.getChildren().get(0));
        String right = generateTerm((LRParser.NonTerminalNode) tail.getChildren().get(1), code);
       // System.out.println("处理右边"+right);
        String temp = newTemp();
        code.add(temp + " = " + left + " " + op + " " + right);

        if (tail.getChildren().size() > 2) {
            return handleAdditiveExpression(temp, (LRParser.NonTerminalNode) tail.getChildren().get(2), code);
        }
        return temp;
    }

    private String generateTerm(LRParser.NonTerminalNode term, List<String> code) {

        String left = generateFactor((LRParser.NonTerminalNode) term.getChildren().get(0), code);
        //System.out.println("处理term："+left);

        if (term.getChildren().size() > 1) {
            LRParser.NonTerminalNode termTail = (LRParser.NonTerminalNode) term.getChildren().get(1);
            return handleMultiplicativeExpression(left, termTail, code);
        }
        return left;
    }

    private String handleMultiplicativeExpression(String left, LRParser.NonTerminalNode tail, List<String> code) {
        if (tail.getChildren().isEmpty()) return left;

        String op = getTokenValue(tail.getChildren().get(0));
        String right = generateFactor((LRParser.NonTerminalNode) tail.getChildren().get(1), code);
        //System.out.println("处理term："+right);
        String temp = newTemp();
        code.add(temp + " = " + left + " " + op + " " + right);

        if (tail.getChildren().size() > 2) {
            return handleMultiplicativeExpression(temp, (LRParser.NonTerminalNode) tail.getChildren().get(2), code);
        }
        return temp;
    }

    private String generateFactor(LRParser.NonTerminalNode factor, List<String> code) {
        LRParser.Node first = factor.getChildren().get(0);

        if (first instanceof LRParser.TokenNode){
        if (isParenthesizedExpression(factor)) {
            //System.out.println("Generating parenthesized expression");
            return generateExpression((LRParser.NonTerminalNode) factor.getChildren().get(1), code);
        }
            return getTokenValue(first);
        }

        else if (first instanceof LRParser.NonTerminalNode) {
            return generateExpression((LRParser.NonTerminalNode) first, code);
        }

        throw new RuntimeException("Invalid factor: " + factor);
    }

    private boolean isParenthesizedExpression(LRParser.NonTerminalNode factor) {
        return factor.getChildren().size() == 3 &&
                getTokenValue(factor.getChildren().get(0)).equals("(") &&
                getTokenValue(factor.getChildren().get(2)).equals(")");
    }
    private String getTokenValue(LRParser.Node node) {
        if (node instanceof LRParser.TokenNode tokenNode) return tokenNode.token.lexeme();
        if (node instanceof LRParser.NonTerminalNode nt && !nt.getChildren().isEmpty())
            return getTokenValue(nt.getChildren().get(0));
        return "?";
    }

    private String newTemp() {
        return "t" + (tempCount++);
    }

    private String newLabel() {
        return "L" + (labelCount++);
    }
}